#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Oracle到DM数据库数据迁移工具
"""

import yaml
import oracledb
import dmPython
import logging
import time
import sys
import os
from typing import List, Dict, Any, Optional

# 设置 oracledb 的客户端编码
oracledb.defaults.encoding = 'UTF-8'


class DatabaseConfig:
    """数据库配置类"""
    def __init__(self, host: str, port: int, service_name: str, username: str, password: str):
        self.host = host
        self.port = port
        self.service_name = service_name
        self.username = username
        self.password = password


class MigrationConfig:
    """迁移配置类"""
    def __init__(self, config_file: str):
        with open(config_file, 'r', encoding='utf-8') as f:
            config = yaml.safe_load(f)
        
        source_db_config = config['source_db']
        self.source_db = DatabaseConfig(
            host=source_db_config['host'],
            port=source_db_config['port'],
            service_name=source_db_config['service_name'],
            username=source_db_config['username'],
            password=source_db_config['password']
        )
        
        target_db_config = config['target_db']
        self.target_db = DatabaseConfig(
            host=target_db_config['host'],
            port=target_db_config['port'],
            service_name=target_db_config['service_name'],
            username=target_db_config['username'],
            password=target_db_config['password']
        )
        
        self.tables_to_migrate = config.get('tables_to_migrate', [])
        self.tables_to_exclude = config.get('tables_to_exclude', [])
        self.batch_size = config.get('batch_size', 1000)
        self.retry_count = config.get('retry_count', 3)


class DatabaseConnection:
    """数据库连接基类"""
    def __init__(self, config: DatabaseConfig):
        self.config = config
        self.connection = None
    
    def connect(self):
        raise NotImplementedError
    
    def disconnect(self):
        if self.connection:
            self.connection.close()
    
    def execute_query(self, query: str, params: Optional[tuple] = None) -> List[Dict]:
        raise NotImplementedError


class OracleConnection(DatabaseConnection):
    """Oracle数据库连接类"""
    def connect(self):
        # 使用 oracledb 3.4.0 版本的新连接方式，直接传递参数
        self.connection = oracledb.connect(
            host=self.config.host,
            port=self.config.port,
            service_name=self.config.service_name,
            user=self.config.username,
            password=self.config.password
        )
        return self.connection
    
    def execute_query(self, query: str, params: Optional[tuple] = None) -> List[Dict]:
        cursor = self.connection.cursor()
        if params:
            cursor.execute(query, params)
        else:
            cursor.execute(query)
        
        columns = [col[0].lower() for col in cursor.description]
        rows = cursor.fetchall()
        # 处理 LOB 类型数据
        processed_rows = []
        for row in rows:
            processed_row = []
            for value in row:
                # 如果是 LOB 类型，读取其内容
                if hasattr(value, 'read'):
                    try:
                        content = value.read()
                        # 处理可能的编码问题
                        if isinstance(content, str):
                            # 处理特殊字符
                            content = content.replace('\xa0', ' ')
                            content = content.replace('\ufffd', ' ')
                        processed_row.append(content)
                    except Exception as e:
                        # 如果读取失败，记录原始值
                        processed_row.append(value)
                else:
                    # 处理普通字符串值中的特殊字符
                    if isinstance(value, str):
                        value = value.replace('\xa0', ' ')
                        value = value.replace('\ufffd', ' ')
                    processed_row.append(value)
            processed_rows.append(processed_row)
        
        result = [dict(zip(columns, row)) for row in processed_rows]
        cursor.close()
        return result


class DMConnection(DatabaseConnection):
    """达梦数据库连接类"""
    def connect(self):
        self.connection = dmPython.connect(
            server=self.config.host,
            port=self.config.port,
            user=self.config.username,
            password=self.config.password,
            schema=self.config.service_name
        )
        return self.connection
    
    def execute_query(self, query: str, params: Optional[tuple] = None) -> List[Dict]:
        cursor = self.connection.cursor()
        if params:
            cursor.execute(query, params)
        else:
            cursor.execute(query)
        
        columns = [col[0].lower() for col in cursor.description]
        rows = cursor.fetchall()
        # 处理 LOB 类型数据
        processed_rows = []
        for row in rows:
            processed_row = []
            for value in row:
                # 如果是 LOB 类型，读取其内容
                if hasattr(value, 'read'):
                    try:
                        content = value.read()
                        # 处理可能的编码问题
                        if isinstance(content, str):
                            # 处理特殊字符
                            content = content.replace('\xa0', ' ')
                            content = content.replace('\ufffd', ' ')
                        processed_row.append(content)
                    except Exception as e:
                        # 如果读取失败，记录原始值
                        processed_row.append(value)
                else:
                    # 处理普通字符串值中的特殊字符
                    if isinstance(value, str):
                        value = value.replace('\xa0', ' ')
                        value = value.replace('\ufffd', ' ')
                    processed_row.append(value)
            processed_rows.append(processed_row)
        
        result = [dict(zip(columns, row)) for row in processed_rows]
        cursor.close()
        return result


class MigrationResult:
    """迁移结果类"""
    def __init__(self, table_name: str):
        self.table_name = table_name
        self.status = "未知"
        self.record_count = 0
        self.error_message = ""
        self.start_time = None
        self.end_time = None
    
    @property
    def duration(self):
        if self.start_time and self.end_time:
            return self.end_time - self.start_time
        return 0


class DataMigrationTool:
    """数据迁移工具主类"""
    def __init__(self, config_file: str):
        self.config = MigrationConfig(config_file)
        self.source_conn = None
        self.target_conn = None
        self.results = []
        self.logger = self._setup_logger()
    
    def _setup_logger(self):
        logger = logging.getLogger("DataMigrationTool")
        logger.setLevel(logging.INFO)
        
        # 创建文件处理器
        file_handler = logging.FileHandler("migration.log", encoding='utf-8')
        file_handler.setLevel(logging.INFO)
        
        # 创建控制台处理器
        console_handler = logging.StreamHandler(sys.stdout)
        console_handler.setLevel(logging.INFO)
        
        # 设置日志格式
        formatter = logging.Formatter(
            '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
        )
        file_handler.setFormatter(formatter)
        console_handler.setFormatter(formatter)
        
        logger.addHandler(file_handler)
        logger.addHandler(console_handler)
        
        return logger
    
    def connect_databases(self):
        """连接源和目标数据库"""
        try:
            # 连接Oracle数据库
            self.source_conn = OracleConnection(self.config.source_db)
            self.source_conn.connect()
            self.logger.info("成功连接到Oracle数据库")
            
            # 连接达梦数据库
            self.target_conn = DMConnection(self.config.target_db)
            self.target_conn.connect()
            self.logger.info("成功连接到达梦数据库")
            
            return True
        except Exception as e:
            self.logger.error(f"数据库连接失败: {str(e)}")
            return False
    
    def disconnect_databases(self):
        """断开数据库连接"""
        if self.source_conn:
            self.source_conn.disconnect()
            self.logger.info("已断开Oracle数据库连接")
        
        if self.target_conn:
            self.target_conn.disconnect()
            self.logger.info("已断达梦数据库连接")
    
    def get_all_tables(self) -> List[str]:
        """获取Oracle数据库中所有表名"""
        query = """
        SELECT table_name 
        FROM user_tables 
        ORDER BY table_name
        """
        result = self.source_conn.execute_query(query)
        return [row['table_name'] for row in result]
    
    def get_table_columns(self, table_name: str) -> List[Dict]:
        """获取表的列信息"""
        query = """
        SELECT column_name, data_type, data_length, data_precision, data_scale, nullable
        FROM user_tab_columns 
        WHERE table_name = :table_name
        ORDER BY column_id
        """
        return self.source_conn.execute_query(query, (table_name,))
    
    def get_table_count(self, table_name: str) -> int:
        """获取表的记录数"""
        query = f"SELECT COUNT(*) as count FROM {table_name}"
        result = self.source_conn.execute_query(query)
        return result[0]['count']
    
    def table_exists_in_target(self, table_name: str) -> bool:
        """检查表是否在目标数据库中存在"""
        query = """
        SELECT table_name 
        FROM user_tables 
        WHERE table_name = :table_name
        """
        result = self.target_conn.execute_query(query, (table_name,))
        return len(result) > 0
    
    def migrate_table_data(self, table_name: str) -> MigrationResult:
        """迁移单个表的数据"""
        result = MigrationResult(table_name)
        result.start_time = time.time()
        
        try:
            self.logger.info(f"开始迁移表 {table_name}")
            
            # 检查目标表是否存在
            if not self.table_exists_in_target(table_name):
                result.status = "忽略"
                result.error_message = f"目标数据库中不存在表 {table_name}"
                self.logger.warning(f"表 {table_name} 在目标数据库中不存在，跳过迁移")
                result.end_time = time.time()
                return result
            
            # 获取表记录数
            record_count = self.get_table_count(table_name)
            if record_count == 0:
                result.status = "成功"
                result.record_count = 0
                self.logger.info(f"表 {table_name} 无数据需要迁移")
                result.end_time = time.time()
                return result
            
            # 获取目标表的列信息，用于处理约束和数据类型
            target_columns = self.get_table_columns(table_name)
            column_info = {col['column_name'].lower(): col for col in target_columns}
            
            # 清空目标表数据
            truncate_query = f"TRUNCATE TABLE {table_name}"
            cursor = self.target_conn.connection.cursor()
            cursor.execute(truncate_query)
            self.target_conn.connection.commit()
            cursor.close()
            
            # 分批迁移数据
            offset = 0
            migrated_count = 0
            
            while offset < record_count:
                # 构造分页查询语句，使用Oracle兼容的语法
                select_query = f"""
                SELECT * FROM (
                    SELECT t.*, ROWNUM rn FROM {table_name} t
                    WHERE ROWNUM <= :end_row
                ) WHERE rn > :start_row
                """
                
                batch_data = self.source_conn.execute_query(
                    select_query, 
                    (offset + self.config.batch_size, offset)
                )
                
                if not batch_data:
                    break
                
                # 过滤掉 rn 列，避免在插入时出现问题
                filtered_batch_data = []
                for row in batch_data:
                    # 移除 rn 列
                    filtered_row = {k: v for k, v in row.items() if k != 'rn'}
                    filtered_batch_data.append(filtered_row)
                
                # 插入到目标数据库
                if filtered_batch_data:
                    # 构造插入语句
                    columns = list(filtered_batch_data[0].keys())
                    placeholders = ", ".join([f":{i+1}" for i in range(len(columns))])
                    insert_query = f"""
                    INSERT INTO {table_name} ({", ".join(columns)}) 
                    VALUES ({placeholders})
                    """
                    
                    # 批量插入
                    cursor = self.target_conn.connection.cursor()
                    for row in filtered_batch_data:
                        # 处理可能包含非法字符的数据
                        values = []
                        for col in columns:
                            value = row[col]
                            col_name = col.lower()
                            
                            # 处理字符串编码问题
                            if isinstance(value, str):
                                try:
                                    # 确保字符串可以被正确编码
                                    value.encode('utf-8')
                                except UnicodeEncodeError:
                                    # 如果无法编码，尝试替换非法字符
                                    value = value.encode('utf-8', errors='replace').decode('utf-8')
                                # 特别处理\xa0字符（不间断空格）
                                value = value.replace('\xa0', ' ')
                                # 特别处理\ufffd字符（替换字符）
                                value = value.replace('\ufffd', ' ')
                                
                                # # 根据目标列信息截断过长的字符串
                                # if col_name in column_info:
                                #     max_length = column_info[col_name].get('data_length', 0)
                                #     if max_length > 0 and len(value) > max_length:
                                #         value = value[:max_length]
                                #         self.logger.warning(f"表 {table_name} 列 {col} 的值被截断至 {max_length} 字符")
                                # elif len(value) > 4000:  # 对未知列也进行限制
                                #     value = value[:4000]
                                #     self.logger.warning(f"表 {table_name} 列 {col} 的值被截断至 4000 字符")
                            # 处理None值，特别是对于非空约束字段
                            elif value is None and col_name in column_info:
                                # 检查是否是非空字段
                                nullable = column_info[col_name].get('nullable', 'Y')
                                data_type = column_info[col_name].get('data_type', '').upper()
                                
                                if nullable == 'N':  # 非空字段
                                    if 'TIME' in data_type or 'DATE' in data_type:
                                        from datetime import datetime
                                        value = datetime.now()
                                        self.logger.warning(f"表 {table_name} 列 {col} 设置默认时间值")
                                    elif 'NUMBER' in data_type or 'INT' in data_type:
                                        value = 0
                                        self.logger.warning(f"表 {table_name} 列 {col} 设置默认数值0")
                                    else:
                                        value = ''
                                        self.logger.warning(f"表 {table_name} 列 {col} 设置默认空字符串")
                            
                            values.append(value)
                        
                        try:
                            cursor.execute(insert_query, values)
                        except Exception as insert_error:
                            # 如果插入失败，记录具体错误并尝试处理
                            error_str = str(insert_error)
                            if "字符串截断" in error_str or "string truncation" in error_str.lower() or "CODE:-70005" in error_str:
                                # 处理字符串截断问题
                                self.logger.warning(f"检测到字符串截断问题，尝试处理...")
                                # 严格限制字符串长度以适应目标字段
                                new_values = []
                                for i, val in enumerate(values):
                                    col_name = columns[i].lower()
                                    if isinstance(val, str) and col_name in column_info:
                                        max_length = column_info[col_name].get('data_length', 4000)  # 默认4000
                                        if max_length > 0 and len(val) > max_length:
                                            val = val[:max_length]
                                            self.logger.warning(f"表 {table_name} 列 {columns[i]} 的值被截断至 {max_length} 字符")
                                    elif isinstance(val, str) and len(val) > 4000:  # 对于未知列也进行限制
                                        val = val[:4000]
                                        self.logger.warning(f"表 {table_name} 列 {columns[i]} 的值被截断至 4000 字符")
                                    new_values.append(val)
                                cursor.execute(insert_query, new_values)
                            elif "非空约束" in error_str or "NOT NULL" in error_str.upper() or "CODE:-6609" in error_str:
                                # 处理非空约束违反问题
                                self.logger.warning(f"检测到非空约束违反，尝试处理...")
                                # 为违反非空约束的字段提供默认值
                                new_values = []
                                for i, val in enumerate(values):
                                    col_name = columns[i].lower()
                                    if val is None and col_name in column_info:
                                        # 检查是否是非空字段
                                        nullable = column_info[col_name].get('nullable', 'Y')
                                        data_type = column_info[col_name].get('data_type', '').upper()
                                        
                                        if nullable == 'N':  # 非空字段
                                            if 'TIME' in data_type or 'DATE' in data_type:
                                                from datetime import datetime
                                                val = datetime.now()
                                                self.logger.warning(f"表 {table_name} 列 {columns[i]} 设置默认时间值")
                                            elif 'NUMBER' in data_type or 'INT' in data_type:
                                                val = 0
                                                self.logger.warning(f"表 {table_name} 列 {columns[i]} 设置默认数值0")
                                            else:
                                                val = ''
                                                self.logger.warning(f"表 {table_name} 列 {columns[i]} 设置默认空字符串")
                                    new_values.append(val)
                                cursor.execute(insert_query, new_values)
                            else:
                                # 重新抛出其他异常
                                raise insert_error
                    
                    self.target_conn.connection.commit()
                    cursor.close()
                
                migrated_count += len(filtered_batch_data)
                offset += self.config.batch_size
                
                # 显示进度
                progress = min(offset, record_count) / record_count * 100
                self.logger.info(f"表 {table_name} 迁移进度: {progress:.1f}% ({min(offset, record_count)}/{record_count})")
            
            result.status = "成功"
            result.record_count = migrated_count
            self.logger.info(f"成功迁移表 {table_name}，共 {migrated_count} 条记录")
            
        except Exception as e:
            result.status = "失败"
            result.error_message = str(e)
            self.logger.error(f"迁移表 {table_name} 失败: {str(e)}")
            # 回滚事务
            try:
                self.target_conn.connection.rollback()
            except:
                pass
        
        result.end_time = time.time()
        return result
    
    def get_tables_to_migrate(self) -> List[str]:
        """获取需要迁移的表列表"""
        # 如果配置了特定表，则只迁移这些表
        if self.config.tables_to_migrate:
            tables = self.config.tables_to_migrate
        else:
            # 否则迁移所有表
            tables = self.get_all_tables()
        
        # 排除指定的表
        if self.config.tables_to_exclude:
            tables = [table for table in tables if table not in self.config.tables_to_exclude]
        
        return tables
    
    def migrate(self) -> List[MigrationResult]:
        """执行数据迁移"""
        start_time = time.time()
        self.logger.info("开始数据迁移任务")
        
        # 获取需要迁移的表
        tables = self.get_tables_to_migrate()
        self.logger.info(f"需要迁移的表: {', '.join(tables)}")
        
        # 迁移每个表
        for table_name in tables:
            try:
                table_result = self.migrate_table_data(table_name)
                self.results.append(table_result)
            except Exception as e:
                # 单表迁移失败不影响整体迁移
                result = MigrationResult(table_name)
                result.status = "失败"
                result.error_message = str(e)
                result.start_time = time.time()
                result.end_time = time.time()
                self.results.append(result)
                self.logger.error(f"迁移表 {table_name} 时发生未处理异常: {str(e)}")
        
        end_time = time.time()
        duration = end_time - start_time
        self.logger.info(f"数据迁移任务完成，总耗时: {duration:.2f} 秒")
        
        return self.results
    
    def generate_report(self) -> Dict[str, Any]:
        """生成迁移报告"""
        total_tables = len(self.results)
        successful_tables = [r for r in self.results if r.status == "成功"]
        failed_tables = [r for r in self.results if r.status == "失败"]
        ignored_tables = [r for r in self.results if r.status == "忽略"]
        
        total_records = sum(r.record_count for r in successful_tables)
        
        report = {
            "summary": {
                "total_tables": total_tables,
                "successful_tables": len(successful_tables),
                "failed_tables": len(failed_tables),
                "ignored_tables": len(ignored_tables),
                "total_records": total_records,
                "success_rate": len(successful_tables) / total_tables * 100 if total_tables > 0 else 0
            },
            "details": [
                {
                    "table_name": r.table_name,
                    "status": r.status,
                    "record_count": r.record_count,
                    "error_message": r.error_message,
                    "duration": round(r.duration, 2)
                }
                for r in self.results
            ]
        }
        
        return report
    
    def print_report(self):
        """打印迁移报告"""
        report = self.generate_report()
        summary = report["summary"]
        details = report["details"]
        
        # 确保输出使用UTF-8编码
        import io
        import sys
        
        # 设置标准输出编码为UTF-8
        if sys.stdout.encoding != 'utf-8':
            try:
                sys.stdout.reconfigure(encoding='utf-8')
            except:
                pass  # 如果无法重新配置，则继续使用当前编码
        
        print("\n========== 数据迁移报告 ==========")
        print(f"总表数: {summary['total_tables']}")
        print(f"成功迁移表数: {summary['successful_tables']}")
        print(f"迁移失败表数: {summary['failed_tables']}")
        print(f"忽略表数: {summary['ignored_tables']}")
        print(f"总记录数: {summary['total_records']}")
        print(f"迁移成功率: {summary['success_rate']:.2f}%")
        print("\n详细信息:")
        print("-" * 80)
        for detail in details:
            # 处理可能的编码问题
            table_name = detail['table_name']
            status = detail['status']
            error_message = detail['error_message']
            
            # 确保字符串可以正确显示
            try:
                table_name.encode('gbk')
            except UnicodeEncodeError:
                table_name = table_name.encode('utf-8', errors='replace').decode('utf-8')
                
            try:
                error_message.encode('gbk')
            except UnicodeEncodeError:
                error_message = error_message.encode('utf-8', errors='replace').decode('utf-8')
            
            # 处理特殊字符
            table_name = table_name.replace('\ufffd', ' ')
            error_message = error_message.replace('\ufffd', ' ')
            
            print(f"表名: {table_name:<20} "
                  f"状态: {status:<4} "
                  f"记录数: {detail['record_count']:<8} "
                  f"耗时: {detail['duration']:<6.2f}s "
                  f"错误: {error_message}")
        print("=" * 80)


def main():
    """主函数"""
    if len(sys.argv) > 1:
        config_file = sys.argv[1]
    else:
        config_file = "config.yaml"
    
    if not os.path.exists(config_file):
        # 处理编码问题
        error_msg = f"配置文件 {config_file} 不存在"
        try:
            print(error_msg)
        except UnicodeEncodeError:
            print(error_msg.encode('utf-8', errors='replace').decode('utf-8'))
        sys.exit(1)
    
    # 创建迁移工具实例
    migration_tool = DataMigrationTool(config_file)
    
    try:
        # 连接数据库
        if not migration_tool.connect_databases():
            sys.exit(1)
        
        # 执行迁移
        migration_tool.migrate()
        
        # 生成并打印报告
        migration_tool.print_report()
        
    except KeyboardInterrupt:
        interrupt_msg = "\n用户中断迁移任务"
        try:
            print(interrupt_msg)
        except UnicodeEncodeError:
            print(interrupt_msg.encode('utf-8', errors='replace').decode('utf-8'))
    except Exception as e:
        error_msg = f"迁移过程中发生错误: {str(e)}"
        try:
            print(error_msg)
        except UnicodeEncodeError:
            print(error_msg.encode('utf-8', errors='replace').decode('utf-8'))
    finally:
        # 断开数据库连接
        migration_tool.disconnect_databases()


if __name__ == "__main__":
    main()